"""基于 ReAct 模式的心理健康聊天 Agent"""

import asyncio
import logging
from datetime import datetime
from typing import Dict, Any, List, Optional, AsyncGenerator
from langchain.agents import create_react_agent, AgentExecutor
from langchain_core.prompts import PromptTemplate
from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.outputs import LLMResult
from langchain_core.messages import BaseMessage
from langchain_core.agents import AgentAction, AgentFinish
from app.core.factories import create_llm_instance
from app.core.tools.react_psychological_tools import (
    analyze_user_intent,
    check_user_safety,
    retrieve_knowledge_documents,
    rerank_documents,
    generate_empathetic_response
)

logger = logging.getLogger(__name__)


class ReActStreamingCallback(BaseCallbackHandler):
    """ReAct Agent 流式回调处理器，捕获思考链和工具调用过程"""
    
    def __init__(self, stream_queue: asyncio.Queue):
        super().__init__()
        self.stream_queue = stream_queue
        self.step_count = 0
        
    async def on_llm_start(self, serialized: Dict[str, Any], prompts: List[str], **kwargs) -> None:
        """LLM 开始时的回调"""
        await self.stream_queue.put({
            "type": "llm_start",
            "message": "AI 开始思考...",
            "timestamp": datetime.now().isoformat()
        })
        
    async def on_llm_end(self, response: LLMResult, **kwargs) -> None:
        """LLM 结束时的回调"""
        if response.generations and response.generations[0]:
            content = response.generations[0][0].text
            
            # 解析 ReAct 格式的输出
            if "Thought:" in content:
                thought_part = content.split("Thought:")[1].split("Action:")[0].strip() if "Action:" in content else content.split("Thought:")[1].strip()
                await self.stream_queue.put({
                    "type": "thought",
                    "content": thought_part,
                    "step": self.step_count,
                    "timestamp": datetime.now().isoformat()
                })
                
            if "Action:" in content and "Action Input:" in content:
                action_part = content.split("Action:")[1].split("Action Input:")[0].strip()
                action_input_part = content.split("Action Input:")[1].strip()
                
                await self.stream_queue.put({
                    "type": "action",
                    "action": action_part,
                    "action_input": action_input_part,
                    "step": self.step_count,
                    "timestamp": datetime.now().isoformat()
                })
                
    async def on_agent_action(self, action: AgentAction, **kwargs) -> None:
        """Agent 执行动作时的回调"""
        self.step_count += 1
        await self.stream_queue.put({
            "type": "action_start",
            "tool": action.tool,
            "tool_input": str(action.tool_input),
            "step": self.step_count,
            "timestamp": datetime.now().isoformat()
        })
        
    async def on_tool_start(self, serialized: Dict[str, Any], input_str: str, **kwargs) -> None:
        """工具开始执行时的回调"""
        tool_name = serialized.get("name", "unknown_tool")
        await self.stream_queue.put({
            "type": "tool_start",
            "tool_name": tool_name,
            "input": input_str,
            "step": self.step_count,
            "timestamp": datetime.now().isoformat()
        })
        
    async def on_tool_end(self, output: str, **kwargs) -> None:
        """工具执行结束时的回调"""
        await self.stream_queue.put({
            "type": "observation",
            "content": output,
            "step": self.step_count,
            "timestamp": datetime.now().isoformat()
        })
        
    async def on_agent_finish(self, finish: AgentFinish, **kwargs) -> None:
        """Agent 完成时的回调"""
        await self.stream_queue.put({
            "type": "final_answer",
            "content": finish.return_values.get("output", ""),
            "timestamp": datetime.now().isoformat()
        })


class PsychologicalReActAgent:
    """基于 ReAct 模式的心理健康聊天 Agent"""
    
    def __init__(self):
        self.llm = create_llm_instance()
        self.tools = [
            analyze_user_intent,
            check_user_safety,
            retrieve_knowledge_documents,
            rerank_documents,
            generate_empathetic_response
        ]
        self.agent = None
        self.agent_executor = None
        self._initialize_agent()
        
    def _initialize_agent(self):
        """初始化 ReAct Agent"""
        # ReAct 提示模板
        react_prompt = PromptTemplate.from_template("""
你是一位专业的AI心理健康助手，使用 ReAct (Reasoning + Acting) 模式来处理用户的心理健康咨询。

你需要按照以下格式进行思考和行动：

Thought: 我需要分析用户的输入，理解他们的需求和情绪状态
Action: analyze_user_intent
Action Input: {{"user_input": "用户输入", "chat_history": ""}}
Observation: [工具返回的结果]

Thought: 基于意图分析结果，我需要检查是否存在安全风险
Action: check_user_safety
Action Input: {{"user_input": "用户输入", "chat_history": ""}}
Observation: [工具返回的结果]

Thought: 如果没有安全风险，我需要检索相关的心理健康知识文档
Action: retrieve_knowledge_documents
Action Input: {{"user_input": "用户输入", "intent": "意图", "max_docs": 5}}
Observation: [工具返回的结果]

Thought: 如果检索到多个文档，我需要对它们进行重排序以提高相关性
Action: rerank_documents
Action Input: {{"user_input": "用户输入", "documents_json": "文档JSON字符串"}}
Observation: [工具返回的结果]

Thought: 现在我有了所有必要的信息，可以生成专业的心理健康回复
Action: generate_empathetic_response
Action Input: {{"user_input": "用户输入", "intent": "意图", "documents_json": "文档JSON字符串", "chat_history": "", "safety_triggered": false}}
Observation: [工具返回的结果]

Thought: 我已经完成了完整的心理健康咨询流程，可以给出最终回复
Final Answer: [基于所有分析和工具结果的最终专业回复]

可用工具：
{tools}

请严格按照 Thought -> Action -> Observation 的格式进行，每次只执行一个动作。

用户输入: {input}
对话历史: {chat_history}

{agent_scratchpad}
""")
        
        # 创建 ReAct Agent
        self.agent = create_react_agent(
            llm=self.llm,
            tools=self.tools,
            prompt=react_prompt
        )
        
        # 创建 Agent 执行器
        self.agent_executor = AgentExecutor(
            agent=self.agent,
            tools=self.tools,
            verbose=True,
            max_iterations=10,
            early_stopping_method="generate",
            handle_parsing_errors=True
        )
        
        logger.info("[PsychologicalReActAgent] ReAct Agent 初始化完成")
        
    async def process_message_stream(
        self,
        user_input: str,
        chat_history: Optional[List[Dict[str, Any]]] = None,
        timeout: int = 60
    ) -> AsyncGenerator[Dict[str, Any], None]:
        """流式处理用户消息，展示完整的思考链过程"""
        start_time = datetime.now()
        chat_history = chat_history or []
        
        # 创建流式队列
        stream_queue = asyncio.Queue()
        callback = ReActStreamingCallback(stream_queue)
        
        try:
            logger.info(f"[PsychologicalReActAgent] 开始 ReAct 流式处理: {user_input[:100]}...")
            
            # 发送开始信号
            yield {
                "type": "react_start",
                "message": "🤖 AI 心理助手开始分析您的问题...",
                "timestamp": start_time.isoformat()
            }
            
            # 准备输入
            agent_input = {
                "input": user_input,
                "chat_history": str(chat_history) if chat_history else "无对话历史"
            }
            
            # 创建异步任务执行 Agent
            async def run_agent():
                try:
                    result = await asyncio.to_thread(
                        self.agent_executor.invoke,
                        agent_input,
                        {"callbacks": [callback]}
                    )
                    await stream_queue.put({"type": "agent_complete", "result": result})
                except Exception as e:
                    await stream_queue.put({"type": "agent_error", "error": str(e)})
                finally:
                    await stream_queue.put({"type": "stream_end"})
            
            # 启动 Agent 任务
            agent_task = asyncio.create_task(run_agent())
            
            # 流式输出结果
            final_result = None
            while True:
                try:
                    # 等待队列中的消息
                    message = await asyncio.wait_for(stream_queue.get(), timeout=1.0)
                    
                    if message["type"] == "stream_end":
                        break
                    elif message["type"] == "agent_complete":
                        final_result = message["result"]
                        continue
                    elif message["type"] == "agent_error":
                        yield {
                            "type": "error",
                            "message": f"处理过程中出现错误: {message['error']}",
                            "timestamp": datetime.now().isoformat()
                        }
                        break
                    else:
                        # 转发所有其他消息
                        yield message
                        
                except asyncio.TimeoutError:
                    # 检查 Agent 任务是否还在运行
                    if agent_task.done():
                        break
                    continue
            
            # 等待 Agent 任务完成
            await agent_task
            
            # 发送最终结果
            execution_time = (datetime.now() - start_time).total_seconds()
            
            if final_result:
                yield {
                    "type": "react_complete",
                    "final_answer": final_result.get("output", ""),
                    "execution_time": execution_time,
                    "timestamp": datetime.now().isoformat()
                }
            else:
                yield {
                    "type": "react_complete",
                    "final_answer": "抱歉，处理过程中出现了问题，请稍后再试。",
                    "execution_time": execution_time,
                    "timestamp": datetime.now().isoformat()
                }
            
            logger.info(f"[PsychologicalReActAgent] ReAct 处理完成，耗时: {execution_time:.2f}秒")
            
        except asyncio.TimeoutError:
            logger.error(f"[PsychologicalReActAgent] ReAct 处理超时 ({timeout}秒)")
            yield {
                "type": "timeout",
                "message": "处理超时，请稍后再试",
                "timestamp": datetime.now().isoformat()
            }
        except Exception as e:
            logger.error(f"[PsychologicalReActAgent] ReAct 处理失败: {e}")
            import traceback
            logger.error(f"[PsychologicalReActAgent] 错误堆栈: {traceback.format_exc()}")
            yield {
                "type": "error",
                "message": f"处理失败: {str(e)}",
                "timestamp": datetime.now().isoformat()
            }
    
    async def process_message(
        self,
        user_input: str,
        chat_history: Optional[List[Dict[str, Any]]] = None,
        timeout: int = 60
    ) -> Dict[str, Any]:
        """非流式处理用户消息"""
        start_time = datetime.now()
        chat_history = chat_history or []
        
        try:
            logger.info(f"[PsychologicalReActAgent] 开始 ReAct 处理: {user_input[:100]}...")
            
            # 准备输入
            agent_input = {
                "input": user_input,
                "chat_history": str(chat_history) if chat_history else "无对话历史"
            }
            
            # 执行 Agent
            result = await asyncio.wait_for(
                asyncio.to_thread(self.agent_executor.invoke, agent_input),
                timeout=timeout
            )
            
            execution_time = (datetime.now() - start_time).total_seconds()
            
            return {
                "response": result.get("output", ""),
                "execution_time": execution_time,
                "success": True,
                "timestamp": datetime.now().isoformat()
            }
            
        except asyncio.TimeoutError:
            logger.error(f"[PsychologicalReActAgent] ReAct 处理超时 ({timeout}秒)")
            raise
        except Exception as e:
            logger.error(f"[PsychologicalReActAgent] ReAct 处理失败: {e}")
            import traceback
            logger.error(f"[PsychologicalReActAgent] 错误堆栈: {traceback.format_exc()}")
            raise


# 创建全局 ReAct Agent 实例
react_agent = PsychologicalReActAgent()